{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "(genai-03-vectordb)=\n",
    "# Experiment tracking with a vector DB\n",
    "\n",
    "This notebook illustrates experiment tracking for document-based models, using the LangChain API to integrate directly with vector databases. You can track documents as artifacts, complete with metadata such as loader type, producer information, and collection details.\n",
    "\n",
    "**In this tutorial**\n",
    "- [SDK reference](#sdk-reference)\n",
    "- [Setup and Imports](#setup-and-imports)\n",
    "- [Milvus configuration](#milvus-configuration)\n",
    "- [Creating an MLRun collection from Milvus](#creating-an-mlrun-collection-from-milvus)\n",
    "- [Create the Milvus vector store](#create-the-milvus-vector-store)\n",
    "- [Working with LangChain Documents and MLRun Artifacts](#working-with-langchain-documents-and-mlrun-artifacts)\n",
    "- [Working with MLRun artifacts in a collection](#working-with-mlrun-artifacts-in-a-collection)\n",
    "- [Using Text Splitters with Artifacts](#using-text-splitters-with-artifacts)\n",
    "- [Using MLRunLoader](#using-mlrunloader)\n",
    "- [Using MLRunLoader with DirectoryLoader](#using-mlrunloader-with-directoryloader)\n",
    "\n",
    "This tutorial uses [Milvus](https://milvus.io/api-reference/pymilvus/v2.4.x/About.md) on a local host for simplicity. To use Milvus without the local host, see [Manage Milvus Connections](https://milvus.io/docs/v2.1.x/manage_connection.md).\n",
    "\n",
    "## SDK reference\n",
    "\n",
    "- [get_vector_store_collection](https://docs.mlrun.org/en/stable/api/mlrun.projects/index.html#mlrun.projects.MlrunProject.get_vector_store_collection)\n",
    "- [log_document](https://docs.mlrun.org/en/stable/api/mlrun.projects/index.html#mlrun.projects.MlrunProject.log_document)\n",
    "- [get_config_profile_attribute](https://docs.mlrun.org/en/stable/api/mlrun.projects/index.html#mlrun.projects.MlrunProject.get_config_profile_attributes)\n",
    "- [MLRunLoader](https://docs.mlrun.org/en/stable/api/mlrun.artifacts/mlrun.artifacts.document.html#mlrun.artifacts.document.MLRunLoader)\n",
    "- [DocumentLoaderSpec](https://docs.mlrun.org/en/stable/api/mlrun.artifacts/mlrun.artifacts.document.html#mlrun.artifacts.document.DocumentLoaderSpec)\n",
    "- [ConfigProfile](https://docs.mlrun.org/en/stable/api/mlrun.datastore/index.html#mlrun.datastore.datastore_profile.ConfigProfile)\n",
    "- [add_artifacts](https://docs.mlrun.org/en/stable/api/mlrun.datastore/index.html#mlrun.datastore.vectorstore.VectorStoreCollection.add_artifacts)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Setup and Imports"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true,
    "tags": []
   },
   "outputs": [],
   "source": [
    "import sys\n",
    "\n",
    "%pip install mlrun langchain langchain-milvus langchain_community backports.tarfile pymilvus[milvus_lite]\n",
    "if sys.version_info.major == 3 and sys.version_info.minor == 9:\n",
    "    %pip install protobuf==3.20.3"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import mlrun\n",
    "import tempfile\n",
    "from langchain_milvus import Milvus\n",
    "from langchain_community.document_loaders import DirectoryLoader\n",
    "from mlrun.artifacts import DocumentLoaderSpec, MLRunLoader\n",
    "from mlrun.datastore.datastore_profile import (\n",
    "    ConfigProfile,\n",
    "    register_temporary_client_datastore_profile,\n",
    ")\n",
    "\n",
    "try:\n",
    "    from langchain.embeddings import FakeEmbeddings\n",
    "    from langchain.text_splitter import CharacterTextSplitter\n",
    "except:\n",
    "    from langchain_core.embeddings import FakeEmbeddings\n",
    "    from langchain_text_splitters import CharacterTextSplitter\n",
    "\n",
    "# Initialize project\n",
    "project = mlrun.get_or_create_project(\n",
    "    name=\"genai-tutorial\", context=\"./\", user_project=True, allow_cross_project=True\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Milvus configuration\n",
    "\n",
    "Create and register a profile representing a Milvus DB. This is done in the project level, once per project.</br>\n",
    "Credentials for the DB may be passed here assuming the code is not introduced into any repo, or they may be provided through project secrets.</br>\n",
    "See [ConfigProfile](https://docs.mlrun.org/en/stable/api/mlrun.datastore/index.html#mlrun.datastore.datastore_profile.ConfigProfile)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "profile = ConfigProfile(\n",
    "    name=\"milvus-config\", public={\"MILVUS_DB\": {\"uri\": \"./milvus_demo.db\"}}\n",
    ")\n",
    "# Register the profile temporarily for the current client session\n",
    "register_temporary_client_datastore_profile(profile)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Creating an MLRun collection from Milvus\n",
    "\n",
    "Create, or use an existing, collection to store the artifact/documents in.</br>\n",
    "Use the configuration stored earlier in the `ConfigProfile` to get the configuration details.</br>\n",
    "You still need to create the actual VectorDB class, since each VectorDB has a different initialization method. \n",
    "See [get_config_profile_attribute](https://docs.mlrun.org/en/stable/api/mlrun.projects/index.html#mlrun.projects.MlrunProject.get_config_profile_attributes)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Initialize embedding model (using FakeEmbeddings for demonstration)\n",
    "embedding_model = FakeEmbeddings(size=3)\n",
    "\n",
    "config = project.get_config_profile_attributes(\"milvus-config\")\n",
    "config"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Create the Milvus vector store\n",
    "\n",
    "In this step you also create the MLRun collection wrapper. See [get_vector_store_collection](https://docs.mlrun.org/en/stable/api/mlrun.projects/index.html#mlrun.projects.MlrunProject.get_vector_store_collection)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "vectorstore = Milvus(\n",
    "    collection_name=\"my_tutorial_collection\",\n",
    "    embedding_function=embedding_model,\n",
    "    connection_args=config[\"MILVUS_DB\"],\n",
    "    auto_id=True,\n",
    ")\n",
    "\n",
    "# Create MLRun collection wrapper\n",
    "collection = project.get_vector_store_collection(vector_store=vectorstore)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Working with LangChain documents and MLRun artifacts"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create a sample document\n",
    "def create_sample_document(content, dir=None):\n",
    "    with tempfile.NamedTemporaryFile(\n",
    "        mode=\"w\", suffix=\".txt\", delete=False, dir=dir\n",
    "    ) as temp_file:\n",
    "        temp_file.write(content)\n",
    "        return temp_file.name\n",
    "\n",
    "\n",
    "# Create and log an MLRun artifact\n",
    "file_path = create_sample_document(\"Sample content for demonstration\")\n",
    "artifact = project.log_document(\"sample-doc\", local_path=file_path)\n",
    "\n",
    "# Convert MLRun artifact to LangChain documents\n",
    "langchain_docs = artifact.to_langchain_documents()\n",
    "print(\"LangChain document content:\", langchain_docs[0].page_content)\n",
    "print(\"LangChain document metadata:\", langchain_docs[0].metadata)\n",
    "\n",
    "# Add LangChain documents to collection\n",
    "milvus_ids = collection.add_documents(langchain_docs)\n",
    "print(\"Documents added with IDs:\", milvus_ids)\n",
    "\n",
    "# Search in collection\n",
    "results = collection.similarity_search(\"sample\", k=1)\n",
    "print(\"Search results:\", [doc.page_content for doc in results])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    " ## Working with MLRun artifacts in a collection"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Add artifacts directly to collection\n",
    "artifact1 = project.log_document(\n",
    "    \"doc1\", local_path=create_sample_document(\"First document\")\n",
    ")\n",
    "artifact2 = project.log_document(\n",
    "    \"doc2\", local_path=create_sample_document(\"Second document\")\n",
    ")\n",
    "\n",
    "# Add multiple artifacts at once\n",
    "milvus_ids = collection.add_artifacts([artifact1, artifact2])\n",
    "print(\"Artifacts added with IDs:\", milvus_ids)\n",
    "\n",
    "# Get back as LangChain documents\n",
    "search_results = collection.similarity_search(\"first\")\n",
    "print(\"Retrieved document:\", search_results[0].page_content)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Using text splitters with Artifacts\n",
    "An artifact represents an original document source file (pdf, doc, etc.). Langchain uses a text splitter to convert such sources into chunks of text called document. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create a text splitter\n",
    "splitter = CharacterTextSplitter(separator=\"\\n\", chunk_size=100, chunk_overlap=20)\n",
    "\n",
    "# Create a longer document\n",
    "long_text = \"This is a longer document.\\n\" * 5\n",
    "long_doc = project.log_document(\n",
    "    \"long-doc\", local_path=create_sample_document(long_text)\n",
    ")\n",
    "\n",
    "# Add artifact with splitting\n",
    "collection_split = project.get_vector_store_collection(\n",
    "    vector_store=Milvus(\n",
    "        collection_name=\"split_collection\",\n",
    "        embedding_function=embedding_model,\n",
    "        connection_args=config[\"MILVUS_DB\"],\n",
    "        auto_id=False,\n",
    "    ),\n",
    ")\n",
    "\n",
    "# Add with custom IDs for chunks\n",
    "ids = collection_split.add_artifacts([long_doc], splitter=splitter, ids=[\"doc1\"])\n",
    "print(\"Generated chunk IDs:\", ids)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Using MLRunLoader\n",
    "MLRunLoader is a wrapper. It receives langchain loader as a parameter (for example `Langchain.PDFloader`, or `Langchain.CSVLoader`), and calls this underlying loader for all its purposes. In addition, it adds the source file as an MLRun artifact.\n",
    "\n",
    "This flow uses [MLRunLoader](https://docs.mlrun.org/en/stable/api/mlrun.artifacts/mlrun.artifacts.document.html#mlrun.artifacts.document.MLRunLoader) for creating instances of a dynamically defined document loader."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create a document loader specification\n",
    "loader_spec = DocumentLoaderSpec(\n",
    "    loader_class_name=\"langchain_community.document_loaders.TextLoader\",\n",
    "    src_name=\"file_path\",\n",
    ")\n",
    "\n",
    "# Create and use MLRunLoader\n",
    "file_path = create_sample_document(\"Content for MLRunLoader test\")\n",
    "loader = MLRunLoader(\n",
    "    source_path=file_path,\n",
    "    loader_spec=loader_spec,\n",
    "    artifact_key=\"loaded-doc\",\n",
    "    producer=project,\n",
    ")\n",
    "\n",
    "# Load documents\n",
    "documents = loader.load()\n",
    "print(\"Loaded document content:\", documents[0].page_content)\n",
    "\n",
    "# Verify artifact creation\n",
    "artifact = project.get_artifact(\"loaded-doc\")\n",
    "print(\"Created artifact key:\", artifact.key)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Using MLRunLoader with DirectoryLoader\n",
    "Langchain DirectoryLoader loads all the files in the directory by calling the langchain loader. When you pass `MLRunLoader`, all the source files are added as MLRun artifacts.\n",
    "See [DocumentLoaderSpec](https://docs.mlrun.org/en/stable/api/mlrun.artifacts/mlrun.artifacts.document.html#mlrun.artifacts.document.DocumentLoaderSpec)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create a directory with multiple documents\n",
    "temp_dir = tempfile.mkdtemp()\n",
    "create_sample_document(\"First file content\", dir=temp_dir)\n",
    "create_sample_document(\"Second file content\", dir=temp_dir)\n",
    "\n",
    "# Configure loader specification\n",
    "artifact_loader_spec = DocumentLoaderSpec(\n",
    "    loader_class_name=\"langchain_community.document_loaders.TextLoader\",\n",
    "    src_name=\"file_path\",\n",
    ")\n",
    "\n",
    "# Create directory loader with MLRunLoader\n",
    "dir_loader = DirectoryLoader(\n",
    "    temp_dir,\n",
    "    glob=\"**/*.*\",\n",
    "    loader_cls=MLRunLoader,\n",
    "    loader_kwargs={\n",
    "        \"loader_spec\": artifact_loader_spec,\n",
    "        \"artifact_key\": \"dir_doc%%\",  # %% will be replaced with unique identifier\n",
    "        \"producer\": project,\n",
    "        \"upload\": False,\n",
    "    },\n",
    ")\n",
    "\n",
    "# Load all documents\n",
    "documents = dir_loader.load()\n",
    "print(f\"Loaded {len(documents)} documents\")\n",
    "\n",
    "# List created artifacts\n",
    "artifacts = project.list_artifacts(kind=\"document\")\n",
    "matching_artifacts = [\n",
    "    art for art in artifacts if art[\"metadata\"][\"key\"].startswith(\"dir_doc\")\n",
    "]\n",
    "\n",
    "print(\"Created artifacts:\", [art[\"metadata\"][\"key\"] for art in matching_artifacts])"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "py311",
   "language": "python",
   "name": "conda-env-.conda-py311-py"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
